{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## U-Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import optim\n",
    "from torchvision.transforms import transforms\n",
    "from torch import nn\n",
    "import torch.utils.data as data\n",
    "from torch.utils.data import DataLoader\n",
    "import PIL.Image as Image\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "import os\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"] = \"TRUE\"\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取数据的路径\n",
    "def make_dataset(root):\n",
    "    imgs = []\n",
    "    # 计算共有多少张原始图片\n",
    "    n = len(os.listdir(root))//2\n",
    "    for i in range(n):\n",
    "        # 找到00i.png的路径\n",
    "        img = os.path.join(root, '%03d.png'%i)\n",
    "        # 找到00i_mask.png的路径\n",
    "        mask = os.path.join(root, '%03d_mask.png'%i)\n",
    "        # 添加至列表\n",
    "        imgs.append((img, mask))\n",
    "    return imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LiverDataset(data.Dataset):\n",
    "    \n",
    "    def __init__(self, root, transform=None, target_transform=None):\n",
    "        imgs = make_dataset(root)\n",
    "        self.imgs = imgs\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        x_path, y_path = self.imgs[index]\n",
    "        img_x = Image.open(x_path)\n",
    "        img_y = Image.open(y_path)\n",
    "        if self.transform is not None:\n",
    "            img_x = self.transform(img_x)\n",
    "        if self.target_transform is not None:\n",
    "            img_y = self.target_transform(img_y)\n",
    "        return img_x, img_y\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_transforms = transforms.Compose([\n",
    "    transforms.ToTensor(), \n",
    "    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n",
    "])\n",
    "# mask只需转为Tensor\n",
    "y_transforms = transforms.ToTensor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 4\n",
    "liver_dataset = LiverDataset('data/liver/train', transform=x_transforms, target_transform=y_transforms)\n",
    "dataloaders = DataLoader(liver_dataset, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 3, 512, 512])\n",
      "torch.Size([4, 1, 512, 512])\n"
     ]
    }
   ],
   "source": [
    "for x, y in dataloaders:\n",
    "    print(x.shape)\n",
    "    print(y.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# U_Net模型中的双卷积网络结构\n",
    "class DoubleConv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(DoubleConv, self).__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            # 此处包含padding，为了使输出图像与输入图像大小相同\n",
    "            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1), \n",
    "            nn.BatchNorm2d(out_ch), \n",
    "            nn.ReLU(inplace=True), \n",
    "                        \n",
    "            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), \n",
    "            nn.BatchNorm2d(out_ch), \n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "    def forward(self, input):\n",
    "        return self.conv(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Unet(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(Unet, self).__init__()\n",
    "        \n",
    "        # 特征图大小不变\n",
    "        self.conv1 = DoubleConv(in_ch, 64)\n",
    "        \n",
    "        # 特征图大小长宽减半\n",
    "        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)\n",
    "        \n",
    "        self.conv2 = DoubleConv(64, 128)\n",
    "        self.pool2 = nn.MaxPool2d(2)\n",
    "        \n",
    "        self.conv3 = DoubleConv(128, 256)\n",
    "        self.pool3 = nn.MaxPool2d(2)\n",
    "        \n",
    "        self.conv4 = DoubleConv(256, 512)\n",
    "        self.pool4 = nn.MaxPool2d(2)\n",
    "        \n",
    "        self.conv5 = DoubleConv(512, 1024)\n",
    "        \n",
    "        \n",
    "        # 长宽翻倍，通道数减半\n",
    "        self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2)\n",
    "        \n",
    "        \n",
    "        self.conv6 = DoubleConv(1024, 512)\n",
    "        self.up7 = nn.ConvTranspose2d(512, 256, 2, stride=2)\n",
    "        self.conv7 = DoubleConv(512, 256)\n",
    "        self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2)\n",
    "        self.conv8 = DoubleConv(256, 128)\n",
    "        self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2)\n",
    "        self.conv9 = DoubleConv(128, 64)\n",
    "        self.conv10 = nn.Conv2d(64, out_ch, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        c1 = self.conv1(x)\n",
    "        p1 = self.pool1(c1)\n",
    "        c2 = self.conv2(p1)\n",
    "        p2 = self.pool2(c2)\n",
    "        c3 = self.conv3(p2)\n",
    "        p3 = self.pool3(c3)\n",
    "        c4 = self.conv4(p3)\n",
    "        p4 = self.pool4(c4)\n",
    "        c5 = self.conv5(p4)\n",
    "        print(c5.shape)\n",
    "        \n",
    "        \n",
    "        up_6 = self.up6(c5)\n",
    "        # 通道维拼接 [N, C, H, W]\n",
    "        merge6 = torch.cat([up_6, c4], dim=1)\n",
    "        c6 = self.conv6(merge6)\n",
    "        up_7 = self.up7(c6)\n",
    "        merge7 = torch.cat([up_7, c3], dim=1)\n",
    "        c7 = self.conv7(merge7)\n",
    "        up_8 = self.up8(c7)\n",
    "        merge8 = torch.cat([up_8, c2], dim=1)\n",
    "        c8 = self.conv8(merge8)\n",
    "        up_9 = self.up9(c8)\n",
    "        merge9 = torch.cat([up_9, c1], dim=1)\n",
    "        c9 = self.conv9(merge9)\n",
    "        c10 = self.conv10(c9)\n",
    "        out = nn.Sigmoid()(c10)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输入图像有3个通道，标签图像有1个通道\n",
    "net = Unet(3, 1).to(device)\n",
    "loss = torch.nn.BCELoss()\n",
    "optimizer = optim.Adam(net.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.randn(4, 3, 512, 512).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 1024, 32, 32])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.4746, 0.5202, 0.5007,  ..., 0.5311, 0.5116, 0.5150],\n",
       "          [0.4633, 0.4943, 0.5558,  ..., 0.5530, 0.5078, 0.4359],\n",
       "          [0.3844, 0.5814, 0.4026,  ..., 0.4806, 0.3109, 0.5004],\n",
       "          ...,\n",
       "          [0.5780, 0.5810, 0.5605,  ..., 0.4301, 0.3645, 0.4728],\n",
       "          [0.4249, 0.5015, 0.5782,  ..., 0.5632, 0.4015, 0.6253],\n",
       "          [0.5205, 0.5249, 0.4500,  ..., 0.5206, 0.3976, 0.5445]]],\n",
       "\n",
       "\n",
       "        [[[0.5136, 0.5717, 0.5534,  ..., 0.5541, 0.4408, 0.5163],\n",
       "          [0.4596, 0.5742, 0.5869,  ..., 0.6626, 0.6493, 0.5478],\n",
       "          [0.3513, 0.5869, 0.4394,  ..., 0.6015, 0.5700, 0.5562],\n",
       "          ...,\n",
       "          [0.5170, 0.5092, 0.6208,  ..., 0.4923, 0.6029, 0.5248],\n",
       "          [0.4656, 0.6171, 0.6039,  ..., 0.6119, 0.3779, 0.5538],\n",
       "          [0.4576, 0.5254, 0.4858,  ..., 0.5856, 0.4313, 0.5676]]],\n",
       "\n",
       "\n",
       "        [[[0.4700, 0.5592, 0.6095,  ..., 0.6347, 0.5079, 0.5056],\n",
       "          [0.4930, 0.4456, 0.5160,  ..., 0.5599, 0.5947, 0.6045],\n",
       "          [0.5320, 0.5545, 0.3734,  ..., 0.5955, 0.4678, 0.5882],\n",
       "          ...,\n",
       "          [0.4300, 0.6422, 0.4932,  ..., 0.6085, 0.4143, 0.4285],\n",
       "          [0.3872, 0.5293, 0.3743,  ..., 0.5552, 0.4212, 0.5145],\n",
       "          [0.4479, 0.5266, 0.4664,  ..., 0.4899, 0.4041, 0.5405]]],\n",
       "\n",
       "\n",
       "        [[[0.5200, 0.5560, 0.5587,  ..., 0.5301, 0.5458, 0.5527],\n",
       "          [0.5346, 0.6370, 0.5861,  ..., 0.5725, 0.4976, 0.4709],\n",
       "          [0.4446, 0.4973, 0.4583,  ..., 0.4920, 0.4630, 0.5358],\n",
       "          ...,\n",
       "          [0.5030, 0.5142, 0.4761,  ..., 0.4749, 0.4272, 0.4920],\n",
       "          [0.4662, 0.4953, 0.5602,  ..., 0.4960, 0.4669, 0.5810],\n",
       "          [0.5316, 0.5507, 0.4288,  ..., 0.5448, 0.4614, 0.4870]]]],\n",
       "       device='cuda:0', grad_fn=<SigmoidBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, loss, optimizer, dataloaders, num_epochs=20):\n",
    "    for epoch in range(num_epochs):\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs-1))\n",
    "        print('-'*10)\n",
    "        dt_size = len(dataloaders.dataset)\n",
    "        epoch_loss = 0\n",
    "        step = 0\n",
    "        for x, y in dataloaders:\n",
    "            step += 1\n",
    "            inputs = x.to(device)\n",
    "            labels = y.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            l = loss(outputs, labels)\n",
    "            l.backward()\n",
    "            optimizer.step()\n",
    "            epoch_loss += l.item()\n",
    "            if step % 200 == 0:\n",
    "                print('%d/%d, train_loss:%0.3f' % (step, (dt_size-1)//dataloaders.batch_size+1, l.item()))\n",
    "        print('epoch %d loss:%0.3f' % (epoch, epoch_loss))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/19\n",
      "----------\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_18032\\3950934349.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtrain_model\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mnet\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mloss\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdataloaders\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_18032\\663071075.py\u001b[0m in \u001b[0;36mtrain_model\u001b[1;34m(model, loss, optimizer, dataloaders, num_epochs)\u001b[0m\n\u001b[0;32m     13\u001b[0m             \u001b[0moutputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     14\u001b[0m             \u001b[0ml\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mloss\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 15\u001b[1;33m             \u001b[0ml\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     16\u001b[0m             \u001b[0moptimizer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     17\u001b[0m             \u001b[0mepoch_loss\u001b[0m \u001b[1;33m+=\u001b[0m \u001b[0ml\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\anaconda3\\lib\\site-packages\\torch\\_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[1;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[0;32m    486\u001b[0m                 \u001b[0minputs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    487\u001b[0m             )\n\u001b[1;32m--> 488\u001b[1;33m         torch.autograd.backward(\n\u001b[0m\u001b[0;32m    489\u001b[0m             \u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    490\u001b[0m         )\n",
      "\u001b[1;32m~\\anaconda3\\lib\\site-packages\\torch\\autograd\\__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[1;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[0;32m    195\u001b[0m     \u001b[1;31m# some Python versions print out the first line of a multi-line function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    196\u001b[0m     \u001b[1;31m# calls in the traceback and some print out the last line\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 197\u001b[1;33m     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[0;32m    198\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    199\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "model = train_model(net, loss, optimizer, dataloaders)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "liver_val = LiverDataset('data/liver/val', transform=x_transforms, target_transform=y_transforms)\n",
    "liver_val = DataLoader(liver_val, batch_size=1)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for i, data in enumerate(dataloaders):\n",
    "        # 左边真实，右边预测\n",
    "        x, z = data\n",
    "        y = model(x.to(device))\n",
    "        img_y = torch.squeeze(y.cpu()).numpy()\n",
    "        plt.subplot(1, 2, 1)\n",
    "        z = torch.squeeze(z).numpy()\n",
    "        plt.imshow(z)\n",
    "        plt.axis('on')\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.imshow(img_y)\n",
    "        plt.axis('on')\n",
    "        plt.pause(0.01)\n",
    "        filename = 'data/liver/predict/' + 'new_%d.png'%i\n",
    "        Image.fromarray((img_y*255).astype('uint8')).convert('L').save(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn.Conv2d()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
